# +----------------------------------------------------------------------
# | ChatWork
# +----------------------------------------------------------------------
# | 这不是一个自由软件,您只能在不用于商业目的的前提下对程序代码进行修改和使用。
# | 任何企业和个人不允许对程序代码以任何形式任何目的再发布,商业使用请获取授权。
# +----------------------------------------------------------------------
# | Author: ChatWork Team <2474369941@qq.com>
# +----------------------------------------------------------------------
from typing import Union
from fastapi import Query
from pydantic import BaseModel, Field


class KnowWholeIn(BaseModel):
    bot: Union[str, None] = Query(default=None, max_length=32, description="智能体编号")


class KnowSearchIn(BaseModel):
    """ 知识库搜索参数 """
    page_no: int = Query(gt=0, default=1, description="当前页码")
    page_size: int = Query(gt=0, le=200, default=15, description="每页条数")
    keyword: Union[str, None] = Query(default=None, description="关键词")
    owner: Union[str, None] = Query(default="", description="拥有者: [all=全部, private=私有, cooperate=协作]")


class KnowDetailIn(BaseModel):
    """ 知识库详情参数 """
    code: str = Query(..., min_length=1, max_length=32, description="知识库编号")


class KnowAddIn(BaseModel):
    """ 知识库新增参数 """
    logo: str = Field(default="", max_length=250, description="知识库图标")
    name: str = Field(..., min_length=1, max_length=32, description="知识库名称")
    intro: str = Field(max_length=800, default="", description="知识库简介")
    vector_model: str = Field(..., min_length=1, pattern=r"\s*\S+:\s*\S+", description="向量模型")
    splits_model: str = Field(..., min_length=1, pattern=r"\s*\S+:\s*\S+", description="拆分模型")

    @classmethod
    def messages(cls):
        return {
            "name.min_length": "请填写知识库名称",
            "logo.max_length": "图标的链接不能超出250个字符",
            "name.max_length": "知识库名称不能超出32个字符",
            "intro.max_length": "知识库简介不能超出800个字符",
            "vector_model.min_length": "请选择向量训练模型",
            "splits_model.min_length": "请选择文件处理模型",
            "vector_model.pattern": "向量训练模型选择异常",
            "splits_model.pattern": "文件处理模型选择异常"
        }

    class Config:
        json_schema_extra = {
            "example": {
                "logo": "https://www.xx.cn/logo.png",
                "name": "知识库名称",
                "intro": "知识库简介",
                "vector_model": "zhipu:embedding-3",
                "splits_model": "zhipu:glm-4-plus"
            }
        }


class KnowEditIn(BaseModel):
    """ 知识库编辑参数 """
    code: str = Field(..., min_length=1, max_length=32, description="知识库编号")
    logo: str = Field(..., min_length=1, max_length=250, description="知识库图标")
    name: str = Field(..., min_length=1, max_length=100, description="知识库名称")
    intro: str = Field(default="", max_length=800, description="知识库简介")
    vector_model: str = Field(..., min_length=1, max_length=200, description="向量模型")
    splits_model: str = Field(..., min_length=1, max_length=200, description="拆分模型")

    @classmethod
    def messages(cls):
        return {
            "name.min_length": "请填写知识库名称",
            "logo.max_length": "图标的链接不能超出250个字符",
            "name.max_length": "知识库名称不能超出32个字符",
            "intro.max_length": "知识库简介不能超出800个字符",
            "vector_model.min_length": "请选择向量训练模型",
            "splits_model.min_length": "请选择文件处理模型",
            "vector_model.pattern": "向量训练模型选择异常",
            "splits_model.pattern": "文件处理模型选择异常"
        }

    class Config:
        json_schema_extra = {
            "example": {
                "code": "wZS3mj23aj6kP",
                "logo": "https://www.xx.cn/logo.png",
                "name": "知识库名称",
                "intro": "知识库简介",
                "vector_model": "zhipu:embedding-3",
                "splits_model": "zhipu:glm-4-plus"
            }
        }


class KnowDeleteIn(BaseModel):
    """ 知识库删除参数 """
    code: str = Field(..., min_length=1, max_length=32, description="知识库编号")

    class Config:
        json_schema_extra = {
            "example": {
                "code": "wZS3mj23aj6kP"
            }
        }


class KnowChangeVectorIn(BaseModel):
    """ 知识库变更向量模型参数 """
    code: str = Field(..., min_length=1, max_length=32, description="知识库编号")
    vector_model: str = Field(..., min_length=1, max_length=200, description="向量模型")

    class Config:
        json_schema_extra = {
            "example": {
                "code": "wZS3mj23aj6kP",
                "vector_model": "zhipu:embedding-3"
            }
        }


"""--------------- Separator ---------------"""


class KnowWholeVo(BaseModel):
    """ 知识库全部VO """
    code: str = Field(description="编号")
    logo: str = Field(description="图标")
    name: str = Field(description="名称")
    vector_name: str = Field(description="模型名称")
    vector_model: str = Field(description="训练模型")
    is_disable: int = Field(description="是否禁用: [0=否, 1=是]")

    class Config:
        json_schema_extra = {
            "example": {
                "code": "wZS3mj23aj6kP",
                "logo": "https://www.xx.cn/logo.png",
                "name": "知识库名称",
                "vector_name": "OpenAI",
                "vector_model": "openai:text-embedding-ada-002",
                "is_disable": 0
            }
        }


class KnowListVo(BaseModel):
    """ 知识库列表VO """
    code: str = Field(description="编号")
    logo: str = Field(description="图标")
    name: str = Field(description="名称")
    intro: str = Field(description="简介")
    power: int = Field(description="拥有的权限")
    is_owner: bool = Field(description="是否私有的")
    vector_model: str = Field(description="向量模型")
    is_disable: int = Field(description="是否禁用: [0=否, 1=是]")
    create_time: str = Field(description="创建时间")
    update_time: str = Field(description="更新时间")
    user: dict = Field(description="用户信息")

    class Config:
        json_schema_extra = {
            "example": {
                "code": "wZS3mj23aj6kP",
                "logo": "https://www.xx.cn/logo.png",
                "name": "知识库名称",
                "intro": "知识库简介",
                "power": 4,
                "is_owner": True,
                "vector_model": "text-embedding-ada-002",
                "create_time": "2023-03-08 21:28:28",
                "update_time": "2023-03-08 21:28:28",
                "user": {
                    "sn": "80965941",
                    "nickname": "wait",
                    "avatar": "https://www.waitadmin.cn/images/avatar.png"
                }
            }
        }


class KnowDetailVo(BaseModel):
    """ 知识库详情VO """
    code: str = Field(description="编号")
    logo: str = Field(description="图标")
    name: str = Field(description="名称")
    intro: str = Field(description="简介")
    owned: bool = Field(description="是否拥有者: [true=拥有者, false=协作者]")
    power: int = Field(description="拥有的权限: [1=查看者, 2=编辑者, 3=管理者, 4=拥有者]")
    vector_name: str = Field(description="向量名称")
    vector_model: str = Field(description="向量模型")
    splits_model: str = Field(description="向量模型")
    is_disable: int = Field(description="是否禁用: [0=否, 1=是]")
    create_time: str = Field(description="创建时间")
    update_time: str = Field(description="更新时间")

    class Config:
        json_schema_extra = {
            "example": {
                "code": "wZS3mj23aj6kP",
                "logo": "https://www.xx.cn/logo.png",
                "name": "知识库名称",
                "intro": "知识库简介",
                "owned": True,
                "power": 4,
                "vector_name": "智谱清言",
                "vector_model": "zhipu:embedding-3",
                "splits_model": "zhipu:glm-4-plus",
                "is_disable": 0,
                "create_time": "2023-03-08 21:28:28",
                "update_time": "2023-03-08 21:28:28"
            }
        }
